#pragma once

#include <thread>
#include <atomic>

extern "C" {
#include <libavcodec/avcodec.h>
#include <libavformat/avformat.h>
}

#include "ThreadSafeQueue.h"

class Decoder {
public:
    Decoder(ThreadSafeQueue<AVPacket*>& packet_queue, ThreadSafeQueue<AVFrame*>& frame_queue, std::atomic<bool>&  demux_done, std::atomic<bool>& decoding_done)
        : packet_queue(packet_queue), frame_queue(frame_queue), quit(false), demux_done(demux_done), decode_done(decoding_done) {}

    bool init(AVCodecParameters* codecpar) {
        const AVCodec* codec = avcodec_find_decoder(codecpar->codec_id);
        codec_ctx = avcodec_alloc_context3(codec);
        avcodec_parameters_to_context(codec_ctx, codecpar);
        return avcodec_open2(codec_ctx, codec, nullptr) >= 0;
    }

    void start() {
        decode_thread = std::thread(&Decoder::decode, this);
    }

    void stop() {
        std::printf("Try to stop decode thread\n");
        std::fflush(stdout);
        quit = true;
        if (decode_thread.joinable()) {
            decode_thread.join();
        }
        std::printf("Decode thread stopped\n");
        std::fflush(stdout);
    }

    ~Decoder() {
        if (codec_ctx) {
            avcodec_free_context(&codec_ctx);
        }
        std::printf("Decoder destroyed\n");
    }

private:
    void decode() {
        while (!quit) {
            if (!packet_queue.empty()) {
                AVPacket* packet = packet_queue.pop();
                if (avcodec_send_packet(codec_ctx, packet) >= 0) {
                    AVFrame* frame = av_frame_alloc();
                    while (avcodec_receive_frame(codec_ctx, frame) >= 0) {
                        frame_queue.push(frame);
                        frame = av_frame_alloc();
                    }
                    av_frame_free(&frame);
                }
                av_packet_free(&packet);
            } 
            else if (demux_done) {
                quit = true;
                break;
            }
            else {
                std::this_thread::sleep_for(std::chrono::milliseconds(10));
            }
        }

        decode_done = true;
        std::printf("Decode Done\n");
        std::fflush(stdout);
    }

    AVCodecContext* codec_ctx = nullptr;
    ThreadSafeQueue<AVPacket*>& packet_queue;
    ThreadSafeQueue<AVFrame*>& frame_queue;
    std::atomic<bool> quit;
    std::atomic<bool>& decode_done;
    std::atomic<bool>& demux_done;
    std::thread decode_thread;
};
